import os
import random
import pandas as pd
from dotenv import load_dotenv
from sqlalchemy import create_engine, text
import backtrader as bt
from datetime import datetime
import numpy as np
from loguru import logger
from strategies.boss_wong_strategy import BossWongStrategy
from strategies.random_print_strategy import RandomPrintStrategy
from strategies.uncle_gen_strategy import UncleGenStrategy
from utils import load_environment_variables, create_database_engine, query_stock_data, preprocess_data
import utils

# 配置loguru日志
logger.add("logs/stock_analysis.log", rotation="1 day", retention="7 days", level="INFO", encoding="utf-8")


def load_stock_data():
    """主函数：加载和预处理股票数据"""
    logger.info("=" * 50)
    logger.info("开始加载股票数据")
    logger.info("=" * 50)
    
    # 1. 加载环境变量
    env_vars = load_environment_variables()
    
    # 2. 创建数据库连接
    engine = create_database_engine(env_vars)
    
    # 3. 查询数据
    df = query_stock_data(engine, env_vars['START_DATE'], env_vars['END_DATE'])
    
    # 4. 预处理数据
    df = preprocess_data(df)
    
    logger.info("=" * 50)
    logger.info("股票数据加载完成")
    logger.info("=" * 50)
    
    return df

# 2. 定义backtrader数据类
class PandasDataEx(bt.feeds.PandasData):
    params = (
        ('datetime', 'trade_date'),
        ('open', 'open'),
        ('high', 'high'),
        ('low', 'low'),
        ('close', 'close'),
        ('volume', 'vol'),
        ('openinterest', None),
    )

# 3. 策略：每天随机选一只股票，打印其close
# RandomPrintStrategy 已移动到 strategies/random_print_strategy.py

if __name__ == "__main__":
    logger.info("程序启动")
    
    # 4. 加载数据
    df = load_stock_data()

    # 5. 按股票分组，添加到backtrader
    logger.info("开始初始化Backtrader引擎...")
    cerebro = bt.Cerebro()
    
    # 获取环境变量
    env_vars = load_environment_variables()
    faction = env_vars['FACTION']

    ts_codes = df['ts_code'].unique()
    target_ts_codes=set(ts_codes[int(env_vars['TSCODE_START']):int(env_vars['TSCODE_END'])])
    
    stock_count = 0
    for ts_code, group in df.groupby('ts_code'):
        if ts_code not in target_ts_codes:
            continue

        # 以一定概率随机保留数据
        if random.random() > faction:
            continue

        if len(group) < env_vars['LEAST_BARS']:
            logger.debug(f"股票 {ts_code} 的K线数量小于 {env_vars['LEAST_BARS']}，跳过")
            continue    

        group = group.sort_values('trade_date')
        data = PandasDataEx(dataname=group, name=ts_code)
        cerebro.adddata(data, name=ts_code)
        # cerebro.resampledata(data, timeframe=bt.TimeFrame.Weeks, compression=1)
        stock_count += 1
        if stock_count % 500 == 0:
            logger.debug(f"已添加 {stock_count} 只股票到引擎")
    
    logger.info(f"Backtrader引擎初始化完成，共添加 {stock_count} 只股票")
    
    logger.debug(f"使用策略: {env_vars['STRATEGY']}")
    if env_vars['STRATEGY'] == "UNCLE_GEN":
        cerebro.addstrategy(UncleGenStrategy, target_trade_date=utils.TRADE_DATE_MAX)
    elif env_vars['STRATEGY'] == "BOSS_WONG":
        cerebro.addstrategy(BossWongStrategy, target_trade_date=utils.TRADE_DATE_MAX)
    else:
        cerebro.addstrategy(RandomPrintStrategy)

    logger.info("开始执行回测策略...")
    cerebro.run()
    logger.info("回测策略执行完成")
    logger.info("程序结束")
